from flask_socketio import SocketIO, emit, join_room, leave_room, disconnect
from flask import request, current_app
from datetime import datetime
import functools
import sqlite3
from db import get_db
import json
import logging

# 创建 SocketIO 实例
socketio = SocketIO(cors_allowed_origins="*")

# 设置日志
logger = logging.getLogger(__name__)

# 存储客户端连接信息
clients = {}

def init_socketio(app):
    """初始化 Socket.IO"""
    # 从配置文件加载 Socket.IO 配置
    socketio_config = app.config.get('SOCKETIO', {})
    
    # 配置 Socket.IO
    socketio.init_app(
        app,
        ping_timeout=socketio_config.get('ping_timeout', 5000),
        ping_interval=socketio_config.get('ping_interval', 25000),
        max_http_buffer_size=socketio_config.get('max_http_buffer_size', 1e8),
        cors_allowed_origins=socketio_config.get('cors_allowed_origins', '*'),
        async_mode='threading'
    )
    
    logger.info('Socket.IO 初始化完成')
    return socketio

def log_websocket_event(event_type, data=None):
    """记录 WebSocket 事件"""
    current_app.logger.info(f'WebSocket 事件: {event_type}')
    if data:
        current_app.logger.info(f'数据: {data}')

def authenticated_only(f):
    @functools.wraps(f)
    def wrapped(*args, **kwargs):
        if not request.args.get('token'):  # 这里可以根据你的认证需求修改
            disconnect()
            return
        return f(*args, **kwargs)
    return wrapped

@socketio.on('connect')
def handle_connect():
    """处理客户端连接"""
    client_id = request.sid
    clients[client_id] = {
        'connected_at': datetime.now().isoformat(),
        'rooms': set()
    }
    logger.info(f'Client {client_id} connected')
    emit('connection_response', {'status': 'connected', 'client_id': client_id})

@socketio.on('disconnect')
def handle_disconnect():
    """处理客户端断开连接"""
    client_id = request.sid
    if client_id in clients:
        logger.info(f'Client {client_id} disconnected')
        del clients[client_id]

@socketio.on('join')
def handle_join(data):
    """处理加入房间"""
    room = data.get('room')
    if room:
        join_room(room)
        clients[request.sid]['rooms'].add(room)
        logger.info(f'Client {request.sid} joined room {room}')
        emit('room_response', {'status': 'joined', 'room': room})

@socketio.on('leave')
def handle_leave(data):
    """处理离开房间"""
    room = data.get('room')
    if room:
        leave_room(room)
        clients[request.sid]['rooms'].discard(room)
        logger.info(f'Client {request.sid} left room {room}')
        emit('room_response', {'status': 'left', 'room': room})

@socketio.on('message')
def handle_message(data):
    """处理消息"""
    try:
        message_type = data.get('type')
        content = data.get('content')
        room = data.get('room')

        if message_type == 'ping':
            emit('pong', {'timestamp': datetime.now().isoformat()})
        elif room:
            # 发送到特定房间
            emit('message', {
                'type': message_type,
                'content': content,
                'sender': request.sid,
                'timestamp': datetime.now().isoformat()
            }, room=room)
        else:
            # 广播消息
            emit('message', {
                'type': message_type,
                'content': content,
                'sender': request.sid,
                'timestamp': datetime.now().isoformat()
            }, broadcast=True)

    except Exception as e:
        logger.error(f'Error handling message: {e}')
        emit('error', {'message': str(e)})

@socketio.on('scan_package')
def handle_package_scan(data):
    """处理包裹扫描"""
    client_id = request.sid
    package_id = data.get('package_id')
    batch_id = data.get('batch_id')
    
    log_websocket_event('scan_package', {
        'client_id': client_id,
        'package_id': package_id,
        'batch_id': batch_id
    })
    
    if not package_id or not batch_id:
        error_msg = 'Missing package_id or batch_id'
        log_websocket_event('scan_error', {
            'client_id': client_id,
            'error': error_msg
        })
        emit('scan_error', {'error': error_msg})
        return
        
    try:
        db = get_db()
        cursor = db.cursor()
        
        # 检查批次是否存在且未关闭
        cursor.execute('SELECT status FROM batches WHERE batch_id = ?', (batch_id,))
        batch = cursor.fetchone()
        if not batch:
            error_msg = '批次不存在'
            log_websocket_event('scan_error', {
                'client_id': client_id,
                'error': error_msg
            })
            emit('scan_error', {'error': error_msg})
            return
        if batch['status'] == 'closed':
            error_msg = '批次已关闭'
            log_websocket_event('scan_error', {
                'client_id': client_id,
                'error': error_msg
            })
            emit('scan_error', {'error': error_msg})
            return
            
        # 获取当前时间
        now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        
        # 检查包裹是否存在
        cursor.execute('''
            SELECT * FROM packages 
            WHERE package_id = ? AND batch_id = ?
        ''', (package_id, batch_id))
        package = cursor.fetchone()
        
        if not package:
            error_msg = '包裹不存在或不属于指定批次'
            log_websocket_event('scan_error', {
                'client_id': client_id,
                'error': error_msg
            })
            emit('scan_error', {'error': error_msg})
            return
            
        # 更新包裹状态
        if package['status'] == 'unscanned':
            cursor.execute('''
                UPDATE packages 
                SET status = 'scanned', 
                    first_scan = ?,
                    last_scan = ?
                WHERE package_id = ? AND batch_id = ?
            ''', (now, now, package_id, batch_id))
            log_websocket_event('package_first_scan', {
                'client_id': client_id,
                'package_id': package_id,
                'batch_id': batch_id
            })
        else:
            cursor.execute('''
                UPDATE packages 
                SET last_scan = ?
                WHERE package_id = ? AND batch_id = ?
            ''', (now, package_id, batch_id))
            log_websocket_event('package_rescan', {
                'client_id': client_id,
                'package_id': package_id,
                'batch_id': batch_id
            })
            
        # 更新批次统计
        cursor.execute('''
            UPDATE batches 
            SET scanned_count = (
                SELECT COUNT(*) FROM packages 
                WHERE batch_id = ? AND status = 'scanned'
            ),
            update_time = ?
            WHERE batch_id = ?
        ''', (batch_id, now, batch_id))
        
        # 记录扫描日志
        cursor.execute('''
            INSERT INTO scan_logs (package_id, scan_time, batch_id, status)
            VALUES (?, ?, ?, ?)
        ''', (package_id, now, batch_id, 'scanned'))
        
        db.commit()
        
        # 获取更新后的包裹和批次信息
        cursor.execute('''
            SELECT p.*, b.total_count, b.scanned_count
            FROM packages p
            JOIN batches b ON p.batch_id = b.batch_id
            WHERE p.package_id = ? AND p.batch_id = ?
        ''', (package_id, batch_id))
        
        updated_package = cursor.fetchone()
        total_count = updated_package['total_count']
        scanned_count = updated_package['scanned_count']
        completion = round(scanned_count / total_count * 100, 1) if total_count > 0 else 0
        
        # 发送扫描结果
        scan_result = {
            'success': True,
            'package_id': package_id,
            'batch_id': batch_id,
            'status': '已扫描',
            'scan_time': now
        }
        log_websocket_event('scan_success', scan_result)
        emit('scan_result', scan_result, room=batch_id)
        
        # 发送批次统计更新
        batch_stats = {
            'total': total_count,
            'scanned': scanned_count,
            'completion': completion
        }
        log_websocket_event('batch_stats_update', batch_stats)
        emit('batch_stats_update', batch_stats, room=batch_id)
        
    except Exception as e:
        error_msg = str(e)
        current_app.logger.error(f'扫描失败: {error_msg}')
        log_websocket_event('scan_error', {
            'client_id': client_id,
            'error': error_msg
        })
        db.rollback()
        emit('scan_error', {'error': error_msg})

@socketio.on('error')
def handle_error(error):
    """处理错误"""
    client_id = request.sid
    log_websocket_event('error', {
        'client_id': client_id,
        'error': str(error)
    })

class WebSocketHandlers:
    def __init__(self, socketio):
        self.socketio = socketio
        self.setup_handlers()

    def setup_handlers(self):
        # Batch Management Events
        @self.socketio.on('join_batch')
        def handle_join_batch(data):
            batch_id = data.get('batch_id')
            if batch_id:
                room = f'batch_room_{batch_id}'
                join_room(room)
                emit('joined_batch', {
                    'status': 'success',
                    'message': f'Joined batch room {batch_id}',
                    'room': room
                })

        @self.socketio.on('create_batch')
        def handle_create_batch(data):
            batch_data = data.get('batch_data')
            try:
                # TODO: Implement batch creation logic
                new_batch = {
                    'id': batch_data.get('id'),
                    'name': batch_data.get('name'),
                    'created_at': datetime.now().isoformat()
                }
                room = f'batch_room_{new_batch["id"]}'
                # Emit to batch list room
                emit('batch_created', new_batch, room='batch_list_room')
                # Emit to specific batch room
                emit('batch_created', new_batch, room=room)
                return {'status': 'success', 'batch': new_batch}
            except Exception as e:
                return {'status': 'error', 'message': str(e)}

        @self.socketio.on('delete_batch')
        def handle_delete_batch(data):
            batch_id = data.get('batch_id')
            try:
                # TODO: Implement batch deletion logic
                room = f'batch_room_{batch_id}'
                # Notify batch list room
                emit('batch_deleted', {'batch_id': batch_id}, room='batch_list_room')
                # Notify specific batch room
                emit('batch_deleted', {'batch_id': batch_id}, room=room)
                return {'status': 'success', 'message': f'Batch {batch_id} deleted'}
            except Exception as e:
                return {'status': 'error', 'message': str(e)}

        # Report Management Events
        @self.socketio.on('join_report')
        def handle_join_report(data):
            report_id = data.get('report_id')
            if report_id:
                room = f'report_room_{report_id}'
                join_room(room)
                emit('joined_report', {
                    'status': 'success',
                    'message': f'Joined report room {report_id}',
                    'room': room
                })

        @self.socketio.on('generate_report')
        def handle_generate_report(data):
            report_data = data.get('report_data')
            try:
                # TODO: Implement report generation logic
                report_id = report_data.get('id')
                room = f'report_room_{report_id}'
                # Send initial progress
                emit('report_progress', {
                    'report_id': report_id,
                    'progress': 0,
                    'status': 'started'
                }, room=room)
                
                # TODO: Actual report generation would happen here
                # For now, just emit completion
                emit('report_complete', {
                    'report_id': report_id,
                    'status': 'completed',
                    'url': f'/reports/{report_id}'
                }, room=room)
                
                return {'status': 'success', 'report_id': report_id}
            except Exception as e:
                emit('report_error', {
                    'report_id': report_id,
                    'error': str(e)
                }, room=room)
                return {'status': 'error', 'message': str(e)}

        # Global Notification Events
        @self.socketio.on('join_global')
        def handle_join_global():
            join_room('global_notification')
            emit('joined_global', {
                'status': 'success',
                'message': 'Joined global notification room'
            })

        @self.socketio.on('disconnect')
        def handle_disconnect():
            # Clean up logic here if needed
            pass

    def broadcast_system_message(self, message):
        """Utility method to broadcast system messages to all connected clients"""
        self.socketio.emit('system_message', {
            'message': message,
            'timestamp': datetime.now().isoformat()
        }, room='global_notification')

    def update_batch_stats(self, batch_id, stats):
        """Utility method to update batch statistics"""
        room = f'batch_room_{batch_id}'
        self.socketio.emit('batch_stats_update', {
            'batch_id': batch_id,
            'stats': stats,
            'timestamp': datetime.now().isoformat()
        }, room=room)

    def update_report_progress(self, report_id, progress, status):
        """Utility method to update report generation progress"""
        room = f'report_room_{report_id}'
        self.socketio.emit('report_progress', {
            'report_id': report_id,
            'progress': progress,
            'status': status,
            'timestamp': datetime.now().isoformat()
        }, room=room) 